import os
import shutil

"""文件均匀分发工具

将当前目录下的所有文件均匀分配到指定数量的子文件夹中。

用法:
    python copy_files.py

交互流程:
    1. 运行后提示输入子文件夹数量
    2. 输入正整数 N
    3. 自动创建 N 个子文件夹 (subfolder_1 到 subfolder_N)
    4. 将文件均匀分配到这些文件夹中

示例:
    > python copy_files.py
    请输入要创建的子文件夹数量: 3
    操作完成。

注意:
    - 只处理当前目录下的文件
    - 不处理子目录中的文件
    - 原始文件保留在原位置
    - 保持文件原始属性
"""

def create_subfolders_and_copy_files(folder_path, num_subfolders):
    # 获取指定文件夹下的所有文件
    files = [f for f in os.listdir(folder_path) if os.path.isfile(os.path.join(folder_path, f))]
    num_files = len(files)
    # 计算每个子文件夹应分配的文件数量
    files_per_folder = num_files // num_subfolders
    remaining_files = num_files % num_subfolders

    current_file_index = 0
    for i in range(num_subfolders):
        # 创建子文件夹
        subfolder_name = os.path.join(folder_path, f"subfolder_{i + 1}")
        os.makedirs(subfolder_name, exist_ok=True)
        # 确定当前子文件夹要复制的文件数量
        num_files_to_copy = files_per_folder
        if i < remaining_files:
            num_files_to_copy += 1
        # 复制文件到子文件夹
        for _ in range(num_files_to_copy):
            if current_file_index < num_files:
                file_to_copy = files[current_file_index]
                source_file_path = os.path.join(folder_path, file_to_copy)
                destination_file_path = os.path.join(subfolder_name, file_to_copy)
                shutil.copy2(source_file_path, destination_file_path)
                current_file_index += 1


if __name__ == "__main__":
    # 获取当前工作目录作为操作文件夹
    folder_path = os.getcwd()
    try:
        # 获取用户输入的子文件夹数量
        num_subfolders = int(input("请输入要创建的子文件夹数量: "))
        if num_subfolders <= 0:
            print("输入的子文件夹数量必须是正整数。")
        else:
            create_subfolders_and_copy_files(folder_path, num_subfolders)
            print("操作完成。")
    except ValueError:
        print("输入无效，请输入一个有效的整数。")